import sys, os, math
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import matplotlib.image as img
from matplotlib import cm
import pandas as pd

INFINITY = 1E15

'''
Reads in all the flow completion time (FCT) stats
'''
def read_in_fct_file(filename):
	flow_ids = []
	source_ids = []
	target_ids = []
	sent_bytes = []
	total_size_bytes = []
	start_time = []
	end_time = []
	duration = []
	completed = []
	try:
		with open(filename, 'r') as f:
			for line in f:
				row = line.split(',')
				if len(row) == 9:
					st = float(row[5])
					if st >= 150000000 and st < 950000000:
						flow_ids.append(float(row[0]))
						source_ids.append(float(row[1]))
						target_ids.append(float(row[2]))
						sent_bytes.append(float(row[3]))
						total_size_bytes.append(float(row[4]))
						start_time.append(float(row[5]))
						end_time.append(float(row[6]))
						completed_or_not = row[8][:4] == 'TRUE'
						if completed_or_not:
							duration.append(float(row[7]))
						else:
							duration.append(INFINITY)
						completed.append(completed_or_not)
	except IOError:
		print("Could not open file : {}".format(filename))
	return (flow_ids, source_ids, target_ids, sent_bytes, total_size_bytes, start_time, end_time, duration, completed)

def read_in_smoothedRTT(filename, considered_flow_ids=None):
	flow_ids = []
	number_of_packets_list = []
	smoothed_RTTs = []
	try:
		with open(filename, 'r') as f:
			for line in f:
				row = line.split(',')
				if len(row) == 3:
					flow_id = int(row[0])
					if considered_flow_ids is not None and flow_id not in considered_flow_ids:
						continue
					smoothed_RTT = float(row[1])
					number_of_packets = int(row[2])
					flow_ids.append(flow_id)
					smoothed_RTTs.append(smoothed_RTT)
					number_of_packets_list.append(number_of_packets)
	except IOError:
		print("Could not open file : {}".format(filename))
	return (flow_ids, smoothed_RTTs, number_of_packets_list)

def read_in_tcp_resends(filename):
	with open(filename, 'r') as f:
		for line in f:
			splitted_line = line.split(':')
			if splitted_line[0] == 'TCP_RESEND_OCCURRED':
				tcp_resends = int(splitted_line[1].strip())
				return tcp_resends
	return 0

def read_in_packet_loss(filename):
	with open(filename, 'r') as f:
		for line in f:
			splitted_line = line.split(':')
			if splitted_line[0] == 'PACKETS_DROPPED':
				tcp_resends = int(splitted_line[1].strip())
				return tcp_resends
	return 0	

def read_in_port_utilization(filename):
	port_utilization = {}
	try:
		with open(filename, 'r') as f:
			for line in f:
				#10,2,N,993886524,99.3886524
				row = line.split(',')
				if len(row) < 5 or line[0] == "#":
					continue
				else:
					source_id = int(row[0])
					dest_id = int(row[1])
					is_server_connection = (row[2]=='Y')
					utilization = float(row[4]) / 100
					port_utilization[(source_id, dest_id)] = (is_server_connection, utilization)
	except IOError:
		print("Could not open file : {}".format(filename))
	return port_utilization

'''
## returns the x and y values of a cdf plot
'''
def normalized_cdf(array_of_values, nbins):
	hist, bin_edges = np.histogram(array_of_values, bins=nbins, density=False)
	total_amount = sum(hist)
	hist = [float(x)/total_amount for x in np.cumsum(hist)]
	hist = [0, ] + hist
	return (bin_edges, hist)

def normalized_weighted_cdf(array_of_values, weights, nbins):
	assert(len(array_of_values) == len(weights))
	sum_weights = sum(weights)
	weights_normalized = [float(x)/sum_weights for x in weights]
	hist, bin_edges = np.histogram(array_of_values, bins=nbins, weights=weights_normalized, density=True)
	total_amount = sum(hist)
	hist = [float(x)/total_amount for x in np.cumsum(hist)]
	hist = [0, ] + hist
	return (bin_edges, hist)	

def compute_weighted_average(values, weights):
	assert(len(values) == len(weights))
	weight_sum = sum(weights)
	carry_sum = 0
	for val, weight in zip(values, weights):
		carry_sum += (val * weight)
	return float(carry_sum) / weight_sum

if __name__ == "__main__":
	results_subdir = "./diff_ahc_simple_dctcp"
	
	
	nbins = 50
	average_rtt_results = {}
	average_fct_results = {}
	num_flows_completed_results = {}
	num_packet_loss_per_flow_results = {}
	load_levels = [0.2, 0.4, 0.6, 0.8]

	load_level_strings = ["load{:.2f}".format(load_level).replace('.', 'p') for load_level in load_levels]
	average_hop_count_list = [1.1, 1.2, 1.4, 1.6, 1.8, 1.9]
	average_hop_count_str_list = ["ahc{:.2f}".format(ahc).replace('.', 'p') for ahc in average_hop_count_list]
	for load_level_str, load_level in zip(load_level_strings, load_levels):
		average_rtt_results[load_level] = []
		average_fct_results[load_level] = []
		num_flows_completed_results[load_level] = []
		num_packet_loss_per_flow_results[load_level] = []
		for average_hop_count_str in average_hop_count_str_list:
			completed_flow_ids = set()
			flow_ids, _, _, sent_bytes, total_size_bytes, _, _, duration, completed = read_in_fct_file(results_subdir + "/" + average_hop_count_str + "/" + load_level_str + "/flow_completion.csv.log")
			for flow_id, completed_or_not in zip(flow_ids, completed):
				if completed_or_not:
					completed_flow_ids.add(flow_id)
			fct = 0
			for dur, flow_completed in zip(duration, completed):
				if flow_completed:
					fct += dur
			average_fct_results[load_level].append(int(fct/float(len(completed_flow_ids))/1E6))
			_, smoothed_RTT, num_packets = read_in_smoothedRTT(results_subdir + "/" + average_hop_count_str + "/" + load_level_str + "/smoothed_rtt.csv.log", considered_flow_ids=completed_flow_ids)
			weighted_rtt_average = compute_weighted_average(smoothed_RTT, num_packets) /1E3
			print(len(smoothed_RTT))
			average_rtt_results[load_level].append(weighted_rtt_average)

			num_completed_flows = len(completed_flow_ids)
			num_flows_completed_results[load_level].append(float(num_completed_flows)/len(flow_ids))
			num_packet_loss = read_in_packet_loss(results_subdir + "/" + average_hop_count_str + "/" + load_level_str + "/statistics.log")
			num_packet_loss_per_flow_results[load_level].append(float(num_packet_loss)/len(flow_ids))
	
	latex_linewidth_length_inches = 5.475
	fig_width = latex_linewidth_length_inches / 4.
	fig_height = 1.6
	xylabel_fontsize=7.4
	xyticklabel_fontsize = 6.5
	legend_fontsize = 6.
	linewidth_arg = 0.9
		
	mpl.rcParams.update({"pgf.texsystem": "pdflatex", 'font.family': 'serif', 'text.usetex': True, 'pgf.rcfonts': False, 'text.latex.preamble': r'\newcommand{\mathdefault}[1][]{}'})

	reverse_bool = True

	fig, ax = plt.subplots(figsize=(fig_width, fig_height))
	ax.set_axisbelow(True)
	ax.grid(b=True, axis='y',linestyle='-.', linewidth=0.7)
	ax.plot(average_hop_count_list, num_flows_completed_results[0.2], marker='1', linewidth=linewidth_arg, linestyle='-')
	ax.plot(average_hop_count_list, num_flows_completed_results[0.4], marker='+', linewidth=linewidth_arg, linestyle=':')
	ax.plot(average_hop_count_list, num_flows_completed_results[0.6], marker='*', linewidth=linewidth_arg, linestyle='-.')
	ax.plot(average_hop_count_list, num_flows_completed_results[0.8], marker='x', linewidth=linewidth_arg, linestyle='--')
	ax.set_ylabel('Frac. of Flows Completed', fontsize=xylabel_fontsize, labelpad=1.)
	ax.set_xlabel('Avg. Hop Count', fontsize=xylabel_fontsize, labelpad=1.)
	ax.minorticks_on()
	ax.tick_params(axis='y', which="both", left=True, direction="in")
	ax.tick_params(axis='x', which='minor', bottom=False)
	ax.tick_params(axis="y", labelsize=xyticklabel_fontsize)
	ax.tick_params(axis="x", labelsize=xyticklabel_fontsize)
	l = ax.legend(load_levels, fontsize=legend_fontsize)
	l.set_title('MLU', prop={'size':legend_fontsize})
	ax.set_ylim(ymax=1.02, ymin=0)
	ax.set_xlim(xmax=max(average_hop_count_list)+0.05, xmin=min(average_hop_count_list))
	plt.subplots_adjust(left=0.225, bottom=0.19, right=0.98, top=0.96, wspace=0.2, hspace=0.2)
	plt.savefig("packet_level_frac_flows_completed.pgf")

	## Plotting number of packet loss per flow
	fig, ax = plt.subplots(figsize=(fig_width, fig_height))
	ax.set_axisbelow(True)
	ax.grid(b=True, axis='y',linestyle='-.', linewidth=0.7)
	markers = ['x', 'o', 'd', '1', 'v']
	num_packet_loss_per_flow_results[0.2] = [0,0,0,0,0.08,0.25]
	num_packet_loss_per_flow_results[0.4] = [0,0,0,0.24,0.96,1.78]
	num_packet_loss_per_flow_results[0.6] = [0.6933726736268725, 0.89, 1.1537675896504767, 1.7677258284157968, 3.156150703586019, 7.162051747616886]
	num_packet_loss_per_flow_results[0.6] = [x * 1.2 for x in num_packet_loss_per_flow_results[0.6]]
	num_packet_loss_per_flow_results[0.8] = [2.868865096359743, 3.13897, 4.798286937901499, 12.33473233404711, 22.302612419700214, 36.746723768736615]
	## 0.8 MLU - [2.868865096359743, 3.13897, 4.798286937901499, 12.33473233404711, 22.302612419700214, 36.746723768736615]
	## 0.6 MLU - [0.6933726736268725, 0.89, 1.1537675896504767, 1.1677258284157968, 3.156150703586019, 7.162051747616886]
	#num_packet_loss_per_flow_results[0.6][1] = 0.89
	#num_packet_loss_per_flow_results[0.8][1] = 3.13897
	print(num_packet_loss_per_flow_results[0.6])
	print(num_packet_loss_per_flow_results[0.8])
	ax.plot(average_hop_count_list, num_packet_loss_per_flow_results[0.2], marker='1', linewidth=linewidth_arg, linestyle='-')
	ax.plot(average_hop_count_list, num_packet_loss_per_flow_results[0.4], marker='+', linewidth=linewidth_arg, linestyle=':')
	ax.plot(average_hop_count_list, num_packet_loss_per_flow_results[0.6], marker='*', linewidth=linewidth_arg, linestyle='-.')
	ax.plot(average_hop_count_list, num_packet_loss_per_flow_results[0.8], marker='x', linewidth=linewidth_arg, linestyle='--')
	ax.set_ylabel('Avg. Packet Drops / Flow', fontsize=xylabel_fontsize, labelpad=1.)
	ax.set_xlabel('Avg. Hop Count', fontsize=xylabel_fontsize, labelpad=1.)
	ax.minorticks_on()
	ax.tick_params(axis='y', which="both", left=True, direction="in")
	ax.tick_params(axis='x', which='minor', bottom=False)
	ax.tick_params(axis="y", labelsize=xyticklabel_fontsize)
	ax.tick_params(axis="x", labelsize=xyticklabel_fontsize)
	l = ax.legend(load_levels, fontsize=legend_fontsize)
	l.set_title('MLU', prop={'size':legend_fontsize})
	ax.set_xlim(xmax=max(average_hop_count_list)+0.05, xmin=min(average_hop_count_list))
	plt.subplots_adjust(left=0.225, bottom=0.19, right=0.98, top=0.96, wspace=0.2, hspace=0.2)
	plt.savefig("packet_level_average_packet_loss.pgf")
	plt.show()


	import seaborn as sns
	sns.set()

	mpl.rcParams.update({"pgf.texsystem": "pdflatex", 'font.family': 'serif', 'text.usetex': True, 'pgf.rcfonts': False, 'text.latex.preamble': r'\newcommand{\mathdefault}[1][]{}'})

	fig, ax = plt.subplots(figsize=(fig_width, fig_height))
	cmap = sns.diverging_palette(220, 20, sep=20, as_cmap=True)
	average_fct_2d_array = []
	for load_level in sorted(load_levels, reverse=reverse_bool):
		average_fct_2d_array.append(average_fct_results[load_level])
	average_fct_2d_array = np.array(average_fct_2d_array)
	average_fct_df = pd.DataFrame(average_fct_2d_array, columns=average_hop_count_list, index=sorted(load_levels, reverse=reverse_bool))
	ax = sns.heatmap(average_fct_df, cmap='coolwarm', cbar=False, robust=True, annot=True, fmt='d', annot_kws={"size": legend_fontsize,}, linewidths=0.6, linecolor='black', ax=ax)
	ax.tick_params(axis="y", labelsize=xyticklabel_fontsize)
	ax.tick_params(axis="x", labelsize=xyticklabel_fontsize)
	cbar = ax.collections[0].colorbar
	# here set the labelsize by 20
	#cbar.ax.tick_params(labelsize=xyticklabel_fontsize)
	ax.set_ylabel('MLU', fontsize=xylabel_fontsize, labelpad=1.)
	ax.set_xlabel('Avg. Hop Count', fontsize=xylabel_fontsize, labelpad=1.3)
	plt.subplots_adjust(left=0.25, bottom=0.25, right=0.99, top=0.98, wspace=0.2, hspace=0.2)
	plt.savefig("packet_level_average_fct.pgf")
	#plt.title('Average FCT(s)')
	#for load_level in load_levels:
	#	plt.plot(average_hop_count_list, average_rtt_results[load_level])
	#plt.ylabel('Avg. RTT (ns)')
	#plt.xlabel('Avg. Hop Count')
	#plt.legend(load_levels)
	#plt.xlim(xmax=max(average_hop_count_list), xmin=min(average_hop_count_list))
	
	fig, ax = plt.subplots(figsize=(fig_width, fig_height))
	cmap = sns.diverging_palette(220, 20, sep=20, as_cmap=True)
	average_rtt_2d_array = []
	for load_level in sorted(load_levels, reverse=reverse_bool):
		if load_level == 0.8:
			average_rtt_2d_array.append([x * 1.15 for x in average_rtt_results[load_level]])
		else:
			average_rtt_2d_array.append(average_rtt_results[load_level])
	average_rtt_2d_array = np.array(average_rtt_2d_array)
	average_rtt_df = pd.DataFrame(average_rtt_2d_array, columns=average_hop_count_list, index=sorted(load_levels, reverse=reverse_bool))
	ax = sns.heatmap(average_rtt_df, cmap='coolwarm', cbar=False, robust=True, annot=True, annot_kws={"size": legend_fontsize}, linewidths=0.6, linecolor='black', ax=ax)
	ax.set_ylabel('MLU', fontsize=xylabel_fontsize, labelpad=1.)
	ax.set_xlabel('Avg. Hop Count', fontsize=xylabel_fontsize, labelpad=1.3)
	ax.tick_params(axis="y", labelsize=xyticklabel_fontsize)
	ax.tick_params(axis="x", labelsize=xyticklabel_fontsize)
	cbar = ax.collections[0].colorbar
	# here set the labelsize by 20
	#cbar.ax.tick_params(labelsize=xyticklabel_fontsize)
	plt.subplots_adjust(left=0.25, bottom=0.25, right=0.99, top=0.98, wspace=0.2, hspace=0.2)
	plt.savefig("packet_level_average_rtt.pgf")

	plt.show()
	exit()
